# -*: coding: utf-8 -*-
""" 
求解器接口
"""

import os

# thirdparty pkgs
import numpy as np
import matplotlib.pyplot as plt

# private pkgs
from algorithm.source_terms.source_interface import SourceInterface
from utils.optmization import error_S
from algorithm.physical_property.tab_loader import TabLoader
from application.boundary_condition import update_boundary
from utils.crdlog import log
from algorithm.equation.momentum import cal_moment_L, cal_moment_gD
from algorithm.equation.pressure_prime import cal_pressure_prime
from algorithm.equation.velocity_correction import correct_V
from algorithm.equation.rho_correction import correct_rho_with_pressure, correct_rho_with_temperature
from algorithm.equation.mass import cal_alpha_g, cal_alpha_L
from algorithm.equation.energy import cal_energy
from application.input_init import input_init
from utils.numeric import *


def cal_source_terms(params, flow_pattern_type, src_interface):
    """
    计算源项
    Args:
        params: dict, 计算需要的输入参数
        physical_property_calculator: 物性计算器
        flow_pattern_type: np.ndarray, 流型类型
    Returns:
        src_interface: SourceInterface, 源项接口
    """
    # 标准化输入
    standard_input = src_interface._input_adapter(params)
    
    # 根据流型类型选择对应的源项计算方法
    source_calculator = src_interface._apply_flow_pattern(standard_input, flow_pattern_type)
    
    # 计算相变项
    src_interface.cal_phase_transition_terms()
    
    # 计算摩擦力项
    src_interface.cal_friction_force(source_calculator)
    
    return src_interface

def update_physical_property(params: dict, physical_property_calculator: TabLoader):
    physical_property_data = physical_property_calculator.get_physical_properties_data(params["p_node"], params["T_node"])
    params["Cp_g_node"] = physical_property_data["gas_heat_capacity"]
    params["Cp_L_node"] = physical_property_data["liquid_heat_capacity"]
    params["rho_L_derivate_p_node"] = physical_property_data["derivative_of_liquid_density_w.r.t_pressure"]
    params["rho_g_derivate_p_node"] = physical_property_data["derivative_of_gas_density_w.r.t_pressure"]
    params["rho_L_derivate_T_node"] = physical_property_data["derivative_of_liquid_density_w.r.t_temperature"]
    params["rho_g_derivate_T_node"] = physical_property_data["derivative_of_gas_density_w.r.t_temperature"]
    
    # 计算声速
    params["c_g_node"] = np.sqrt(1 / params["rho_g_derivate_p_node"])
    params["c_L_node"] = np.sqrt(1 / params["rho_L_derivate_p_node"])
    
    # 气相质量分数
    params["gas_mass_fraction_node"] = physical_property_data["gas_mass_fraction"]


def update_init_var_from_last_time_step(params):
    """ 从上一时刻的计算结果更新当前时刻的初始值
    
    Args:
        params: dict, 计算需要的输入参数
    """
    params["rho_g_node_iteration_init"] = np.copy(
        params["rho_g_node"]
    )
    params["rho_L_node_iteration_init"] = np.copy(
        params["rho_L_node"]
    )
    params["alpha_g_node_iteration_init"] = np.copy(
        params["alpha_g_node"]
    )
    params["alpha_L_node_iteration_init"] = np.copy(
        params["alpha_L_node"]
    )
    params["alpha_D_node_iteration_init"] = np.copy(
        params["alpha_D_node"]
    )
    params["V_L_segment_iteration_init"] = np.copy(params["V_L_segment"])
    params["V_g_segment_iteration_init"] = np.copy(params["V_g_segment"])
    params["V_D_segment_iteration_init"] = np.copy(params["V_D_segment"])
    params["p_node_iteration_init"] = np.copy(
        params["p_node"])
    params["T_node_iteration_init"] = np.copy(
        params["T_node"])
    params["Cp_g_node_iteration_init"] = np.copy(params["Cp_g_node"])
    params["Cp_L_node_iteration_init"] = np.copy(params["Cp_L_node"])

def update_cfl(params: dict):
    """ 更新CFL数
    """
    delta_V_g_segment = params["V_g_segment"] - params["V_g_segment_iteration_init"]
    delta_V_L_segment = params["V_L_segment"] - params["V_L_segment_iteration_init"]
    delta_V_D_segment = params["V_D_segment"] - params["V_D_segment_iteration_init"]
    max_delta_velocity = max(np.abs(delta_V_g_segment).max(), np.abs(delta_V_L_segment).max(), np.abs(delta_V_D_segment).max())
    
    cfl_values = [0.01, 0.02, 0.03, 0.04, 0.05, 0.06, 0.07, 0.08, 0.09, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9]
    max_delta_velocity_values = [1e-2, 8e-3, 6e-3, 4e-3, 2e-3, 1e-3, 8e-4, 6e-4, 4e-4, 2e-4, 1e-4, 8e-5, 6e-5, 4e-5, 2e-5, 1e-5, 8e-6, 6e-6, 4e-6, 2e-6, 1e-6]
    for cfl, delta_velocity in zip(cfl_values, max_delta_velocity_values):
        if max_delta_velocity > delta_velocity:
            params["cfl"] = cfl
            break
    else:
        params["cfl"] = 1
    return

def cal_dt_by_cfl(params):
    """ 根据CFL数计算时间步长

    Args:
        params: dict, 计算需要的输入参数
    """
    max_velocity = np.maximum(np.abs(params["V_L_segment"]).max(), np.abs(params["V_g_segment"]).max())
    params["dt"] = params["cfl"] * params["dx"] / max_velocity
    if params["dt"] < params["dt_min"]:
        params["dt"] = params["dt_min"]
    if params['dt'] > params["dt_max"]:
        params['dt'] = params["dt_max"] 

def cal_init_field(params: dict, physical_property_calculator: TabLoader, flow_pattern_existence: bool, src_interface: SourceInterface) -> dict:
    """ 计算初始场
    """
    mixture_velocity = params["V_L_segment"]*grid_interp(params["alpha_L_node"], 1) + params["V_g_segment"]*grid_interp(params["alpha_g_node"], 1)
    mean_velocity = np.abs(mixture_velocity).mean()
    L = np.sqrt((params["X_node"][-1] - params["X_node"][0])**2 + (params["h_node"][-1] - params["h_node"][0])**2)
    time_flow_to_end = 1.2*(L / mean_velocity)

    init_calculation_time = 0
    log.info(f"开始计算初始场")
    log.info('===============================================')
    while init_calculation_time < time_flow_to_end:
        # 单个时间步迭代

        single_time_step(params, physical_property_calculator, flow_pattern_existence, src_interface)

        # 更新CFL数
        update_cfl(params)
        
        # 计算时间步长
        cal_dt_by_cfl(params)
        
        # 记录时间
        init_calculation_time += params["dt"]
        
    log.info(f"初始场计算完成")
    log.info('===============================================')



def single_time_step(params: dict, physical_property_calculator: TabLoader,  flow_pattern_existence: bool, src_interface: SourceInterface) -> dict:
    """ 单个时间步内的空间迭代

    Args:
        params: dict 计算需要的输入参数
        rt: ReadTab 读表类实例
        flow_pattern_existence: bool 是否存在流型
        src_interface: SourceInterface 源项接口

    Returns:
        params: dict 经过单个时间步计算更新后的计算参数
    """
    # 开始仿真计算
    Error = 1
    step_count = 0  # 记录迭代步次数
    
    # step 1 保存上一时刻的计算结果为当前时刻的初始值
    update_init_var_from_last_time_step(params)

    # step 2 判断流型 TODO: 写死为分层流
    flow_pattern_type = np.ones(params["segment_number"])  

    # step 3 计算源项
    src_interface = cal_source_terms(params, flow_pattern_type, src_interface)
    gas_wall_friction_force, liquid_wall_friction_force, gas_liquid_friction_force = src_interface.get_friction_force()
    liquid_vapour_mass_transfer_rate, V_i_segment, transition_velocity_from_liquid_film, transition_velocity_from_droplet,liquid_to_droplet_mass_transfer_rate, droplet_to_liquid_mass_transfer_rate, droplet_velocity = src_interface.get_phase_transition_terms()

    while Error > 1e-3:
        # step 4 动量方程
        params["V_L_segment"] = cal_moment_L(params,
        liquid_wall_friction_force, gas_liquid_friction_force, liquid_vapour_mass_transfer_rate, V_i_segment, transition_velocity_from_liquid_film, liquid_to_droplet_mass_transfer_rate, droplet_to_liquid_mass_transfer_rate, droplet_velocity)

        params["V_g_segment"] = cal_moment_gD(params, gas_wall_friction_force, gas_liquid_friction_force, liquid_vapour_mass_transfer_rate, V_i_segment, transition_velocity_from_droplet, liquid_to_droplet_mass_transfer_rate, droplet_to_liquid_mass_transfer_rate, droplet_velocity)

        # step 5 压力修正方程
        p_prime_node = cal_pressure_prime(params, liquid_vapour_mass_transfer_rate)
        params["p_node"] += p_prime_node

        # step 6 速度修正方程
        params["V_g_segment"], params["V_L_segment"] = correct_V(params, p_prime_node)
        
        # step 7 由于(p)变化，第一次更新密度
        params["rho_g_node"], params["rho_L_node"] = correct_rho_with_pressure(params, p_prime_node)
        
        # step 8 能量方程
        params["T_node"] = cal_energy(params, gas_wall_friction_force, liquid_wall_friction_force, gas_liquid_friction_force)
    
        # step 9 由于(T)变化，第二次更新密度
        params["rho_g_node"], params["rho_L_node"] = correct_rho_with_temperature(params, params["T_node"] - params["T_node_iteration_init"])
        
        # step 10 由于(p, T)变化，更新物性
        update_physical_property(params, physical_property_calculator)
        
        # step 11 计算截面分数
        params["alpha_g_node"] = cal_alpha_g(params, liquid_vapour_mass_transfer_rate)
        params["alpha_L_node"] = cal_alpha_L(params, liquid_vapour_mass_transfer_rate, liquid_to_droplet_mass_transfer_rate, droplet_to_liquid_mass_transfer_rate)
        
        # step 11.1 截面分数归一化
        params["alpha_g_node"] = params["alpha_g_node"] / (params["alpha_g_node"] + params["alpha_L_node"])
        params["alpha_L_node"] = params["alpha_L_node"] / (params["alpha_g_node"] + params["alpha_L_node"])

        # step 12 边界条件更新
        update_boundary(params, (params["gas_mass_fraction_node"][0] + params["gas_mass_fraction_node"][1])/2)

        # step 13 迭代目标，压力的修正比
        Error = error_S(p_prime_node) / \
            error_S(params["p_node"])
        
        step_count += 1  # 记录迭代步次数
    
    params["iteration_num"] = step_count
    return params

# 接口函数
def time_process(input: dict, specific_end_time: float = None, flow_pattern_existence: bool = False):
    """ 模拟器时间推进计算
    
    Args:
        input (dict): 输入的参数字典
        specific_end_time (float): 指定计算终止时间(用作调试输出)，默认使用输入的最大时间
        flow_pattern_existence (bool): 是否存在流型
    """
    # 用户输入初始化
    params, physical_property_calculator = input_init(input)

    # 实例化源项接口
    src_interface = SourceInterface()
    
    if specific_end_time is None:
        specific_end_time = params["max_time"]
        
    # 计算初始场
    cal_init_field(params, physical_property_calculator, flow_pattern_existence, src_interface)
    
    count = 1
    while params["record_time"] < specific_end_time:

        # 单个时间步迭代
        single_time_step(params, physical_property_calculator, flow_pattern_existence, src_interface)

        # 更新CFL数
        update_cfl(params)
        
        # 计算时间步长
        cal_dt_by_cfl(params)
        
        # 记录模拟时间
        params["record_time"] += params["dt"]

        # 日志输出
        if params["record_time"] % 1000 <= params["dt"]:
            log.info(
            f'计算到第{params["record_time"]:.4f}s, 当前空间迭代步数为{params["iteration_num"]}, 当前dt为{params["dt"]}'
        )
        count += 1
    
    # TODO:物性修正持液率, 这里做一次修正即可, 目的是为了修正输出结果的持液率
    params["alpha_L_node"] = (1/params["gas_mass_fraction_node"] - 1)*params["rho_g_node"]*params["alpha_g_node"]*grid_extrap(params["V_g_segment"], 1)/(params["rho_L_node"]*grid_extrap(params["V_L_segment"], 1))
    params["alpha_L_node"] = apply_global_quadratic_constraint(params["alpha_L_node"], 0)
    return params

if __name__ == "__main__":
    # 测试用
    from utils.constant import BoundaryConditionType
    from utils.common import get_enum_name_by_value
    from tests.visual_olga_horizontal import load_olga_groundtruth_result_se
    from utils.optmization import error_MSE
    from utils.base_path import get_base_path
    from tests.input.olga_case_data import *
    from tests.input.game_case_data import *

    # 开始测试
    case = game_test_5
    
    params = time_process(input=case, specific_end_time=3600, flow_pattern_existence=False)
    
    # =========================================================
    # 测试和保存数据        
    # =========================================================
    classname = get_enum_name_by_value(BoundaryConditionType, params["boundary_class"])

    savepath = os.path.join(get_base_path(), "data/", classname)
    
    if params["boundary_class"] == BoundaryConditionType.inlet_M_outlet_p.value:

        # olga管道案例可视化
        if case['gt_path']:
            gtpath = case['gt_path']
            olga_gt_path = os.path.join(get_base_path(), gtpath)
            (gt_flow_g, gt_flow_L,
            gt_p, gt_T, gt_alpha_L,
            flow_g, flow_L,
            p, T, alpha_L) = load_olga_groundtruth_result_se(params, olga_gt_path)

            flow_g_error = error_MSE(gt_flow_g, flow_g)
            flow_L_error = error_MSE(gt_flow_L, flow_L)
            p_error = error_MSE(gt_p, p)
            T_error = error_MSE(gt_T, T)
            alpha_L_error = error_MSE(gt_alpha_L, alpha_L)
            mean_gt_flow_g = np.mean(gt_flow_g)  # olga对比值百分比
            mean_gt_flow_L = np.mean(gt_flow_L)  # olga对比值百分比
            mean_gt_p = np.mean(gt_p)  # olga对比值百分比
            mean_gt_T = np.mean(gt_T)  # olga对比值百分比
            mean_gt_alpha_L = np.mean(gt_alpha_L)  # olga对比值百分比

            #画图显示,转译代码时，可删除。
            from tests.plot_olga_horizontal import vis_ploter
            # current_param_name = casename.split(".")[0]
            casename = gtpath.split("/")[-1].split(".")[0]
            is_plot = True
            savepath = os.path.join(get_base_path(), "tests/test_algorithm/data/output/", casename)
            if not os.path.exists(savepath):
                os.makedirs(savepath)
            xalL = np.zeros_like(flow_g)
            xalL[:] = params['X_node'][1:-1]

            is_dynamic_range = False # 是否规范取值范围
            vis_ploter(xalL, gt_flow_g, flow_g,
                        gt_flow_L, flow_L,
                        gt_p, p,
                        gt_T, T,
                        gt_alpha_L, alpha_L,
                        [params["record_time"]],
                        None,
                        is_plot,
                       is_dynamic_range,
                        savepath)
